from roboticstoolbox import DHRobot, RevoluteDH
import roboticstoolbox as rtb
import numpy as np
from spatialmath import SE3

class Jaka(DHRobot):
    def __init__(self):
        super().__init__(
            [
                RevoluteDH(alpha=np.pi/2, d=120.15, name="Joint 1"),
                RevoluteDH(a=-430, name="Joint 2"),
                RevoluteDH(a=-368.5, name="Joint 3"),
                RevoluteDH(alpha=np.pi/2, d=113.5, name="Joint 4"),
                RevoluteDH(alpha=-np.pi/2, d=113.5, name="Joint 5"),
                RevoluteDH(d=107, name="Joint 6")
            ],
            name="Jaka Robot"
        )

def get_object_pose():
    # 目标物体在仿真器坐标系下的位姿矩阵
    T_sim_obj = np.array([
        [1.0000, 0.0000, 0.0000, -700.0000],
        [0.0000, 1.0000, 0.0000, -125.0000],
        [0.0000, 0.0000, 1.0000, 125.0000],
        [0, 0, 0, 1]
    ])
    return SE3(T_sim_obj)

def get_staight_trajectory(T_start,T_end,robot,start_q=[0,0,0,0,0,0]):
    cartisanean_trajectory = rtb.ctraj(T_start, T_end, 10)
    viapoints = []
    last_q = start_q
    for index, matrix in enumerate(cartisanean_trajectory):
        q= robot.ikine_LM(matrix,q0=last_q)
        last_q = q.q 
        viapoints.append(q.q)  
    viapoints = np.array(viapoints)
    trajectory = rtb.mstraj(viapoints , dt=0.1, tacc=0.2, qdmax=5)
    return trajectory








def get_trajectory(T_sim_obj):
    robot = Jaka()
    T_plan_end_initial = robot.fkine(robot.q)
    T_sim_end_initial = SE3(np.array([
        [1, 0, 0, -898.5],
        [0, 0, -1, -221],
        [0, 1, 0, 608.2],
        [0, 0, 0, 1]
    ]))
    T_sim_to_plan = T_plan_end_initial * T_sim_end_initial.inv()
    T_sim_obj = get_object_pose()
    T_obj_to_precatch = SE3(np.array([
        [1, 0, 0, 1.5],
        [0, -1, 0, -246.1],
        [0, 0, -1, 583.2],
        [0, 0, 0, 1]
    ]))
    T_sim_end_target = T_obj_to_precatch * T_sim_obj
    T_sim_initial_to_target = T_sim_end_target * T_sim_end_initial.inv()
    T_plan_end_target = T_sim_to_plan * T_sim_initial_to_target * T_sim_to_plan.inv() * T_plan_end_initial

    T_plan_end_target = SE3.Trans(0,0,50)*T_plan_end_target

    qf = robot.ik_GN(T_plan_end_target)
    trajectory = rtb.jtraj(robot.q, qf[0], t=50)

    T_catch_end_target = SE3.Trans(0,0,-150)*T_plan_end_target

    trajectory1 = get_staight_trajectory(T_plan_end_target,T_catch_end_target,robot,qf[0])

    T_move = SE3.Rx(-180, unit='deg')*SE3.Trans(0,0,200)*T_catch_end_target
    
    qf = robot.ik_GN(T_move)
    trajectory2 = rtb.jtraj(trajectory1.q[-1], qf[0], t=50)


    return trajectory , trajectory1 , trajectory2


